package com.pengxh.app.gb28181.util;

import android.media.Image;

import java.nio.ByteBuffer;

public class ImageConverter {

    public static byte[] convertToNV12(Image image) {
        if (image == null) {
            return null;
        }

        int width = image.getWidth();
        int height = image.getHeight();

        Image.Plane yPlane = image.getPlanes()[0];
        Image.Plane uPlane = image.getPlanes()[1];
        Image.Plane vPlane = image.getPlanes()[2];

        ByteBuffer yBuffer = yPlane.getBuffer();
        ByteBuffer uBuffer = uPlane.getBuffer();
        ByteBuffer vBuffer = vPlane.getBuffer();

        int ySize = width * height;
        int uvSize = ySize / 2; // NV12 中 UV 总大小 = Y 大小 / 2
        byte[] nv12 = new byte[ySize + uvSize];

        // Step 1: Copy Y
        yBuffer.get(nv12, 0, ySize);

        // Step 2: Interleave U and V into UV plane (U first, then V)
        int uvWidth = width / 2;
        int uvHeight = height / 2;

        // 获取 stride（行字节对齐，可能 > width）
        int yStride = yPlane.getRowStride();
        int uStride = uPlane.getRowStride();
        int vStride = vPlane.getRowStride();

        // 如果 stride == width 且无 padding，可直接 bulk copy（优化路径）
        if (uStride == uvWidth && vStride == uvWidth && uPlane.getPixelStride() == 1 && vPlane.getPixelStride() == 1) {
            // Fast path: no padding
            for (int i = 0; i < uvSize / 2; i++) {
                nv12[ySize + i * 2] = uBuffer.get(i);     // U
                nv12[ySize + i * 2 + 1] = vBuffer.get(i); // V
            }
        } else {
            // Safe path: handle padding and pixel stride
            int uvIndex = ySize;
            for (int row = 0; row < uvHeight; row++) {
                for (int col = 0; col < uvWidth; col++) {
                    // U
                    int uOffset = row * uStride + col * uPlane.getPixelStride();
                    // V
                    int vOffset = row * vStride + col * vPlane.getPixelStride();

                    nv12[uvIndex++] = uBuffer.get(uOffset);
                    nv12[uvIndex++] = vBuffer.get(vOffset);
                }
            }
        }
        return nv12;
    }

    public static byte[] convertToNV21(Image image) {
        if (image == null) {
            return null;
        }

        int width = image.getWidth();
        int height = image.getHeight();

        // 获取YUV_420_888格式的三个平面
        Image.Plane[] planes = image.getPlanes();

        // 计算NV21格式所需的字节数
        int ySize = width * height;
        int uvSize = width * height / 4;
        byte[] nv21 = new byte[ySize + uvSize * 2];

        // 复制Y平面数据
        ByteBuffer yBuffer = planes[0].getBuffer();
        int yStride = planes[0].getRowStride();
        int yPixelStride = planes[0].getPixelStride();
        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                nv21[i * width + j] = yBuffer.get(i * yStride + j * yPixelStride);
            }
        }

        // 复制UV平面数据（NV21格式中UV交替排列）
        ByteBuffer uBuffer = planes[1].getBuffer();
        ByteBuffer vBuffer = planes[2].getBuffer();
        int uStride = planes[1].getRowStride();
        int vStride = planes[2].getRowStride();
        int uPixelStride = planes[1].getPixelStride();
        int vPixelStride = planes[2].getPixelStride();

        int uvIndex = ySize;
        for (int i = 0; i < height / 2; i++) {
            for (int j = 0; j < width / 2; j++) {
                // NV21格式是VU交替排列
                nv21[uvIndex++] = vBuffer.get(i * vStride + j * vPixelStride);
                nv21[uvIndex++] = uBuffer.get(i * uStride + j * uPixelStride);
            }
        }

        return nv21;
    }
}
